{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copyright 2019 Google LLC\n",
    "# \n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "#     https://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a target=\"_blank\" href=\"https://colab.research.google.com/github/GoogleCloudPlatform/keras-idiomatic-programmer/blob/master/notebooks/building_in_processing_in_graph.ipynb\">\n",
    "<img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TF 2.0 Transition - Building in Pre-processing into the Model\n",
    "\n",
    "TF 2.0 adds a lot of new features and more powerful representation. This notebook will demonstrate some of the newer features to build (custom) input pre-processing into the graph. What's the benefit to this:\n",
    "\n",
    "    1. Since it is part of the model, one does not need to re-implement the preprocessing on the inference \n",
    "       side.\n",
    "    2. Since it will be added as graph ops, the preprocessing will happen on the GPU (instead of upstream on \n",
    "       CPU) and be faster.\n",
    "    3. The preprocessing graph operations can be optimized by the Tensorflow compiler.\n",
    "    \n",
    "## Objective\n",
    "\n",
    "We will be using the following TF 2.0 features / recommendations:\n",
    "\n",
    "    1. [Recommentation] Use tf.keras for the model building.\n",
    "    2. [Recommendation] Put preprocessing into the model.\n",
    "    3. [Feature] Use @tf.function decorator to convert the Python code for preprocessing into graph ops.\n",
    "    4. [Feature] Use subclassing of layers to define a new layer for the preprocessing.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports\n",
    "\n",
    "If you haven't already, you need to install TF 2.0 beta. If you are running this notebook in colab (which is 1.13 as of this writing), you will need to install TF 2.0 via a cell, as in below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# If not already installed\n",
    "%pip install tensorflow==2.0.0-beta1 "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's import what we will use in this demonstration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "from tensorflow.keras import Model, Input, layers\n",
    "from tensorflow.keras.layers import Flatten, Dense\n",
    "\n",
    "# expected output: 2.0.0-beta1\n",
    "print(tf.__version__)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Layer Subclassing\n",
    "\n",
    "We will start by subclassing the tf.keras layers class to make a new layer type, as follows:\n",
    "\n",
    "    1. Will take an input vector whose shape is specified at instantiation.\n",
    "    2. Will normalize the data between 0 and 1 (assumes pixel data between 0 .. 255).\n",
    "    3. Outputs the normalized input.\n",
    "    4. Has no trainable parameters.\n",
    "    \n",
    "Let's start by showing a basic template for subclassing layers and then explain it:\n",
    "\n",
    "```python\n",
    "class NewLayer(layers.Layer):\n",
    "    def __init__(self):\n",
    "        super(NewLayer, self).__init__()\n",
    "        self.my_vars = blash, blah\n",
    "    \n",
    "    def build(self, input_shape):\n",
    "        \"\"\" Handler for building the layer \"\"\"\n",
    "        self.kernel = blah, blah\n",
    "    \n",
    "    def call(self, inputs):\n",
    "        \"\"\" Handler for layer object as callable \"\"\"\n",
    "        outputs = do something with inputs\n",
    "        return outputs\n",
    "```\n",
    "\n",
    "### Subclassing\n",
    "\n",
    "The first line in the above template `class NewLayer(layers.Layer)` indicates we want to create a new class object named `NewLayer` which is subclassed (derived) from the tf.keras `layers` class. This will give us a custom layer definition.\n",
    "\n",
    "### __init__() method\n",
    "\n",
    "This is the initializer (constructor) for the class object instantiation. We use the initializer to initialize layer specific variables.\n",
    "\n",
    "### build() method\n",
    "\n",
    "This method handles the building of the layer when the model is compiled. A typical action is to define the shape of the kernel (trainable parameters) and initialization of the kernel.\n",
    "\n",
    "### call() method\n",
    "\n",
    "This method handles calling the layer as a callable (function call) for execution in the graph."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Subclassing Our Custom Layer\n",
    "\n",
    "In the code below, we subclass a custom layer for doing preprocessing of the input, and where the preprocessing is converted to graph operations in the model.\n",
    "\n",
    "The first line in the code `class Normalize(layers.Layer)` indicates we want to create a new class object named `Normalize` which is subclassed (derived) from the tf.keras `layers` class. \n",
    "\n",
    "\n",
    "### __init__() method\n",
    "\n",
    "Since we won't have any constants or variables to preserve, we don't have any need to add anything to this method.\n",
    "\n",
    "### build() method\n",
    "\n",
    "Our custom layer won't have any trainable parameters. We will tell the compile process to not set up any gradient descent updates on the kernel during training by setting the `layers` class variable `self.kernel` to `None`.\n",
    "\n",
    "### call() method\n",
    "\n",
    "This is where we add our preprocessing. The parameter `inputs` is the input tensor to the layer during training and prediction. A TF tensor object implements polymorphism to overload operators. We use the overloaded division operator, which will broadcast the division operation across the entire tensor --thus each element will be divided by 255.0.\n",
    "\n",
    "Finally, we add the decorator `@tf.function` to tell **TensorFlow AutoGraph** to convert convert the Python code in this method to graph operations in the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Normalize(layers.Layer):\n",
    "    \"\"\" Custom Layer for Preprocessing Input \"\"\"\n",
    "    def __init__(self):\n",
    "        \"\"\" Constructor \"\"\"\n",
    "        super(Normalize, self).__init__()\n",
    "    \n",
    "    def build(self, input_shape):\n",
    "        \"\"\" Handler for Input Shape \"\"\"\n",
    "        self.kernel = None\n",
    "    \n",
    "    @tf.function\n",
    "    def call(self, inputs):\n",
    "        \"\"\" Handler for layer object is callable \"\"\"\n",
    "        inputs = inputs / 255.0\n",
    "        return inputs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build the Model\n",
    "\n",
    "Let's build a model to train on the MNIST dataset. We will keep it really basic:\n",
    "\n",
    "    1. Use the Functional API method for defining the model.\n",
    "    2. Make the first layer of our model the custom preprocessing layer.\n",
    "    3. The remaining layers are a basic DNN for MNIST."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the input vector for 28x28 MNIST images\n",
    "inputs = Input((28, 28))\n",
    "\n",
    "# The first layer is the preprocessing layer, which is bound to the input vector\n",
    "x = Normalize()(inputs)\n",
    "\n",
    "# Next layer, we flatten the preprocessed input into a 1D vector\n",
    "x = Flatten()(x)\n",
    "\n",
    "# Create a hidden dense layer of 128 nodes\n",
    "x = Dense(128, activation='relu')(x)\n",
    "\n",
    "# Create an output layer for classifying the 10 digits\n",
    "outputs = Dense(10, activation='sigmoid')(x)\n",
    "\n",
    "# Instantiate the model\n",
    "model = Model(inputs, outputs)\n",
    "\n",
    "# Compile the model\n",
    "model.compile(loss='sparse_categorical_crossentropy', optimizer='adam', metrics=['acc'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get the Dataset\n",
    "\n",
    "We will get the tf.keras builtin dataset for MNIST. The dataset is pre-split into train and test data. The data is separated into numpy multi-dimensional arrays for images and labels. The image data is not preprocessed --i.e., all the values are between 0 and 255. The label data is not one-hot-encoded --hence why we compiled with `loss='sparse_categorical_crossentropy'`\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.keras.datasets import mnist\n",
    "\n",
    "# Load the train and test data into memory\n",
    "(x_train, y_train), (x_test, y_test) = mnist.load_data()\n",
    "\n",
    "# Expected output: (60000, 28, 28) , (60000,)\n",
    "print(x_train.shape, y_train.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train the Model\n",
    "\n",
    "Let's now train the model (with the preprocessing built into the model graph) on the unpreprocessed MNIST data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.fit(x_train, y_train, epochs=10, batch_size=32, validation_split=0.1, verbose=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluate the Model\n",
    "\n",
    "Let's now evaluate (prediction) using unpreprocessed test examples on the trained model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "acc = model.evaluate(x_test, y_test)\n",
    "print(acc)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
